package jsr292.weaver;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.UndeclaredThrowableException;
import java.util.Collection;

public class Capabilities {
    private static final Method DEFINE_CLASS;
    private static final Method ADD_TRANSFORMER_1_6; 
    private static final Method RETRANSFORM_CLASSES;
    private static final Object UNSAFE;
    private static final Method UNSAFE_STATIC_FIELD_BASE;
    private static final Method UNSAFE_STATIC_FIELD_OFFSET;
    private static final Method UNSAFE_OBJECT_FIELD_OFFSET;
    static {
        Method m;
        try {
            m = ClassLoader.class.getDeclaredMethod("defineClass", String.class, byte[].class, int.class, int.class);
            m.setAccessible(true);
        } catch (SecurityException e) {
            m = null;
        } catch (NoSuchMethodException e) {
            throw new AssertionError(e);
        }
        DEFINE_CLASS = m;

        
        Method addTransformer, retransformClasses;
        try {
            addTransformer = Instrumentation.class.getMethod("addTransformer", ClassFileTransformer.class, boolean.class);
            retransformClasses = Instrumentation.class.getMethod("retransformClasses", Class[].class);
        } catch (NoSuchMethodException e) {
            addTransformer = null;
            retransformClasses = null;
        }
        ADD_TRANSFORMER_1_6 = addTransformer;
        RETRANSFORM_CLASSES = retransformClasses;
        
        Object unsafe;
        Method staticFieldBase, staticFieldOffset, objectFieldOffset;
        try {
            Class<?> unsafeClass = Class.forName("sun.misc.Unsafe");
            Field unsafeField = unsafeClass.getDeclaredField("theUnsafe");
            unsafeField.setAccessible(true);
            unsafe = unsafeField.get(null);
            
            staticFieldBase = unsafeClass.getMethod("staticFieldBase", Field.class);
            staticFieldOffset = unsafeClass.getMethod("staticFieldOffset", Field.class);
            objectFieldOffset = unsafeClass.getMethod("objectFieldOffset", Field.class);
        } catch(Exception e) {
            unsafe = null;
            staticFieldBase = null;
            staticFieldOffset = null;
            objectFieldOffset = null;
        }
        UNSAFE = unsafe;
        UNSAFE_STATIC_FIELD_BASE = staticFieldBase;
        UNSAFE_STATIC_FIELD_OFFSET = staticFieldOffset;
        UNSAFE_OBJECT_FIELD_OFFSET = objectFieldOffset;
    }
    
    public static boolean canDefine() {
        return DEFINE_CLASS != null;
    }
    
    public static boolean canRetransform() {
        return RETRANSFORM_CLASSES != null;
    }
    
    public static boolean canUseUnsafe() {
        return UNSAFE != null;
    }
    
    public static Class<?> define(ClassLoader classloader, String className, byte[] bytecode) {
        try {
            return (Class<?>)DEFINE_CLASS.invoke(classloader, className, bytecode, 0, bytecode.length);
        } catch (IllegalArgumentException e) {
            throw newAssertionError(e);
        } catch (IllegalAccessException e) {
            throw newAssertionError(e);
        } catch (InvocationTargetException e) {
            throw rethrow(e);
        }
    }
    
    public static void addTransformer(Instrumentation instrumentation, ClassFileTransformer transformer, boolean canRetransform) {
        try {
            ADD_TRANSFORMER_1_6.invoke(instrumentation, transformer, canRetransform);
        } catch (IllegalAccessException e) {
            throw newAssertionError(e);
        } catch (InvocationTargetException e) {
            throw rethrow(e);
        }
    }
    
    public static void retransform(Collection<? extends Class<?>> classes) {
        try {
            RETRANSFORM_CLASSES.invoke(Agent.instrumentation,
                    new Object[]{ classes.toArray(new Class<?>[classes.size()]) });
        } catch (IllegalAccessException e) {
            throw newAssertionError(e);
        } catch (InvocationTargetException e) {
            throw rethrow(e);
        }
    }
    
    // called directly from bytecode, see OptimizerSnippets.unsafe*
    public static Object getUnsafe() {
        return UNSAFE;
    }
    
    public static Object staticFieldBase(Field field) {
        try {
            return UNSAFE_STATIC_FIELD_BASE.invoke(UNSAFE, field);
        } catch (IllegalAccessException e) {
            throw newAssertionError(e);
        } catch (InvocationTargetException e) {
            throw rethrow(e);
        }
    }
    
    // called directly from bytecode, see OptimizerBoundMap
    public static long staticFieldOffset(Field field) {
        try {
            return (Long)UNSAFE_STATIC_FIELD_OFFSET.invoke(UNSAFE, field);
        } catch (IllegalAccessException e) {
            throw newAssertionError(e);
        } catch (InvocationTargetException e) {
            throw rethrow(e);
        }
    }
    
    public static long objectFieldOffset(Field field) {
        try {
            return (Long)UNSAFE_OBJECT_FIELD_OFFSET.invoke(UNSAFE, field);
        } catch (IllegalAccessException e) {
            throw newAssertionError(e);
        } catch (InvocationTargetException e) {
            throw rethrow(e);
        }
    }
    
    private static AssertionError newAssertionError(Exception e) {
        return new AssertionError(e);
    }
    
    private static RuntimeException rethrow(InvocationTargetException e) {
        Throwable cause = e.getCause();
        if (cause instanceof RuntimeException) {
            return (RuntimeException)cause;
        }
        if (cause instanceof Error) {
            throw (Error)cause;
        }
        return new UndeclaredThrowableException(cause);   
    }
}
